import numpy as np
import casadi as ca

from src.utils.geometry_utils import polygon_H_from_vertices
from src.model import Model


class ConvexLocalPlanner:
    def __init__(self, model: Model, timestep: float, Q: np.ndarray, R: np.ndarray, N=20):
        super().__init__()

        self._N = N

        self._model = model
        robot_poly = np.array([
            [0.4, 0.3],
            [-0.4, 0.3],
            [-0.4, -0.3],
            [0.4, -0.3],
        ])
        A_ro, b_ro = polygon_H_from_vertices(robot_poly)

        self._opti = ca.Opti()
        self._opt_controls = self._opti.variable(self._N, self._model.num_control)

        self._opt_states = self._opti.variable(self._N + 1, self._model.num_state)

        self._opt_x0 = self._opti.parameter(self._model.num_state)
        self._opt_xs = self._opti.parameter(self._model.num_state)

        self._obstacle_num = 2
        self._obstacle_point_num = 4
        self._A_os = [self._opti.parameter(self._obstacle_point_num, 2) for _ in range(self._obstacle_num)]
        self._b_os = [self._opti.parameter(self._obstacle_point_num, 1) for _ in range(self._obstacle_num)]

        for i in range(self._N):
            x_next = self._opt_states[i, :] + self._model.casadi_derivatives(self._opt_states[i, :],
                                                                             self._opt_controls[i, :]).T * timestep
            self._opti.subject_to(self._opt_states[i + 1, :] == x_next)

        self._slack_obs = 1e5
        self._slacks = []

        self._obs_slacks = self._opti.variable(self._N, self._obstacle_num)
        for i in range(N):
            for j in range(self._obstacle_num):
                self._opti.subject_to(self._obs_slacks[i, j] > 0)

        def rotation(theta):
            st = ca.sin(theta)
            ct = ca.cos(theta)
            return ca.blockcat([[ct, -st], [st, ct]])

        for i in range(self._N):
            rot_i = rotation(self._opt_states[i, 2])
            traj_i = ca.vertcat(self._opt_states[i, 0], self._opt_states[i, 1])

            for j in range(2):
                lambda_ri_j = self._opti.variable(b_ro.shape[0])
                lambda_oi_j = self._opti.variable(self._b_os[j].shape[0])
                self._opti.subject_to(lambda_ri_j[:] > 0)
                self._opti.subject_to(lambda_oi_j[:] > 0)

                Aot_minus_bo = ca.mtimes(self._A_os[j], traj_i) - self._b_os[j]
                equ1 = - ca.mtimes(b_ro.T, lambda_ri_j) + ca.mtimes(Aot_minus_bo.T, lambda_oi_j)
                equ2 = ca.mtimes(A_ro.T, lambda_ri_j) + ca.mtimes(ca.mtimes(rot_i.T, self._A_os[j].T), lambda_oi_j)
                equ3 = ca.sumsqr(ca.mtimes(self._A_os[j].T, lambda_oi_j))

                self._opti.subject_to(equ1 >= -self._obs_slacks[i, j])
                self._opti.subject_to(equ2 == 0)
                self._opti.subject_to(equ3 == 1)

        self._Q = Q.copy()
        self._R = R.copy()

        self._cost = 0

        for i in range(self._model.limit.shape[0]):
            self._opti.subject_to(
                self._opti.bounded(self._model.limit[i, 0], self._opt_controls[:, i], self._model.limit[i, 1]))

        self._opts_setting = {
            'ipopt.max_iter': 1000,
            'ipopt.print_level': 0,
            'print_time': 0,
            'ipopt.acceptable_tol': 1e-8,
            'ipopt.acceptable_obj_change_tol': 1e-6
        }
        self._opti.solver('ipopt', self._opts_setting)

        self._u0 = np.zeros((N, self._model.num_control))
        self._next_states = np.zeros((N + 1, self._model.num_state))

    def init(self):
        self._opti.subject_to(self._opt_states[0, :] == self._opt_x0.T)

        for i in range(self._N):
            self._cost += ca.mtimes(
                [(self._opt_states[i, :] - self._opt_xs.T), self._Q,
                 (self._opt_states[i, :] - self._opt_xs.T).T]) + ca.mtimes(
                [self._opt_controls[i, :], self._R, self._opt_controls[i, :].T])
        self._cost += self._slack_obs * ca.sumsqr(self._obs_slacks)
        self._opti.minimize(self._cost)

    def compute(self, current_state, desired_state, obstacles):
        self._opti.set_value(self._opt_x0, current_state)
        self._opti.set_value(self._opt_xs, desired_state)

        for i, obstacle in enumerate(obstacles):
            A_oi, b_oi = polygon_H_from_vertices(obstacle)
            self._opti.set_value(self._A_os[i], A_oi)
            self._opti.set_value(self._b_os[i], b_oi)

        self._opti.set_initial(self._opt_controls, self._u0)
        self._opti.set_initial(self._opt_states, self._next_states)

        sol = self._opti.solve()

        u_res = sol.value(self._opt_controls)
        next_states_pred = sol.value(self._opt_states)

        self._u0[:, :] = np.concatenate((u_res[1:], u_res[-1:]))
        self._next_states[:, :] = np.concatenate((next_states_pred[1:], next_states_pred[-1:]))

        return u_res[0]

    @property
    def next_states(self):
        return self._next_states.copy()
